// ebpf/prefetch_kern.c (版本 4 - 验证器友好版)
#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>

char LICENSE[] SEC("license") = "GPL";

struct event {
    __u32 pid;
    __u64 fault_addr;
    long stride;
};

#define MAX_HISTORY 8
#define STABILITY_THRESHOLD 3

struct access_history {
    __u64 addrs[MAX_HISTORY];
    long last_stride;
    __u32 stable_count;
    __u32 cursor;
};

struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 256 * 1024);
} rb SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 10240);
    __type(key, __u32);
    __type(value, struct access_history);
} history_map SEC(".maps");

SEC("tp/exceptions/page_fault_user")
int handle_page_fault(struct pt_regs *ctx) {
    __u64 fault_addr = PT_REGS_PARM1(ctx);
    __u32 pid = bpf_get_current_pid_tgid() >> 32;

    struct access_history *history;
    struct access_history new_history = {}; // 先在栈上准备好

    history = bpf_map_lookup_elem(&history_map, &pid);
    if (!history) {
        // 如果找不到，我们不能直接返回，因为后面可能会用到 history 指针。
        // 一个安全的模式是：如果找不到，就用栈上的 new_history 的地址。
        // 但更简单且验证器友好的模式是：先插入一个空的，再查找一次。
        bpf_map_update_elem(&history_map, &pid, &new_history, BPF_NOEXIST);
        history = bpf_map_lookup_elem(&history_map, &pid);
        if (!history) {
            // 如果这次还找不到，说明 map 满了或者出错了，直接返回。
            return 0; 
        }
    }

    // --- 关键修改 ---
    // 在使用 history 指针之前，再次进行一个看似多余但对验证器很重要的 NULL 检查。
    // 这会强制验证器在后续代码中都认为 history 是一个有效的指针。
    if (history == NULL) {
        return 0;
    }

    // 计算步长
    __u32 cursor = history->cursor;
    __u32 prev_idx = (cursor - 1 + MAX_HISTORY) % MAX_HISTORY;
    __u64 prev_addr = history->addrs[prev_idx];
    
    // 我们必须检查 cursor 是否越界，即使我们知道它不会
    // 这也是为了满足验证器
    if (cursor >= MAX_HISTORY) {
        cursor = 0;
    }
    
    long current_stride = 0;
    if (prev_addr > 0) {
       current_stride = fault_addr - prev_addr;
    }

    // 更新本次访问记录
    history->addrs[cursor] = fault_addr;
    history->cursor = (cursor + 1) % MAX_HISTORY;
    
    // 过滤噪声
    if (current_stride <= 0 || (current_stride & 4095) != 0) {
        history->stable_count = 0;
        history->last_stride = current_stride;
        return 0;
    }

    // 轻量级模式识别
    if (current_stride == history->last_stride) {
        history->stable_count++;
    } else {
        history->last_stride = current_stride;
        history->stable_count = 1;
    }

    // 达到阈值，发送事件
    if (history->stable_count >= STABILITY_THRESHOLD) {
        struct event *e = bpf_ringbuf_reserve(&rb, sizeof(*e), 0);
        if (e) {
            e->pid = pid;
            e->fault_addr = fault_addr;
            e->stride = history->last_stride;
            bpf_ringbuf_submit(e, 0);
        }
        history->stable_count = 0; // 重置
    }

    return 0;
}
